from engine import *
import random
class Neuron:
    '''
    神经元定义：
    nin： 支持多少个信号的输入 也就是多少个维度的输入
    w： 权重， 也就是每个维度的权重
    b： 偏置， 也就是每个维度的偏置
    只有一个输出
    '''
    def __init__(self, nin):
        self.w = [Value(random.uniform(-1, 1)) for _ in range(nin)]
        self.b = Value(random.uniform(-1, 1))

    def __call__(self, x):
        #w * x + b
        act = sum ((wi * xi for wi, xi in zip(self.w, x)), self.b)
        out = act.tanh()
        return out;
    def parameters(self):
        return self.w + [self.b]
    
class Layer:
    '''
    层定义：
    nout: 支持多少个神经元，每个神经元都有一个输出，也就是说，这个层有多少个输出
    nin: 每个神经元支持多少个维度的输入
    '''
    def __init__(self, nin, nout):
        self.neurons = [Neuron(nin) for _ in range(nout)]

    def __call__(self, x):
        out = [n(x) for n in self.neurons]
        return out[0] if len(out) == 1 else out
    def parameters(self):
        return [p for neuron in self.neurons for p in neuron.parameters()]
    
class MLP:
    '''
    多层感知机定义：
    nin: 支持多少个维度的输入
    nouts: 支持多少个层， 每个层支持多少个神经元，是一个列表
    '''
    def __init__(self, nin, nouts):
        sz = [nin] + nouts
        self.layers = [Layer(sz[i], sz[i+1]) for i in range(len(nouts))]
    
    def __call__(self, x):
        for layer in self.layers:
            x = layer(x)
        return x
    
    def parameters(self):
        return [p for layer in self.layers for p in layer.parameters()]

    
if __name__ == "__main__":
    x = [2.0, 3.0, -1.0]
    xs = [
        [2.0, 3.0, -1.0],
        [3.0, -1.0, 0.5],
        [0.5, 1.0, 1.0],
        [1.0, 1.0, -1.0]
    ]
    ys = [1.0, -1.0, -1.0, 1.0]
    n = MLP(3, [4, 4, 1])
    for k in range(100):
        # forward
        ypred = [n(x) for x in xs]
        loss = sum([(yout - ygt) ** 2 for ygt, yout in zip (ys, ypred)])
        for p in n.parameters():
            p.grad = 0
        # backward
        loss.backward()
        # update
        for p in n.parameters():
            p.data += -0.05 * p.grad
        print(loss)
        